import sys
import os
import os

os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'  # 解决多线程库冲突的警告
import requests

import torch
import numpy as np

import matplotlib.pyplot as plt
from PIL import Image

sys.path.append('..')  # 将上级目录添加到系统路径中，方便导入自定义模块
import models_mae  # 导入模型定义（假设该模块包含 MAE 模型）

# 定义 ImageNet 数据集的均值和标准差，用于图像预处理
imagenet_mean = np.array([0.485, 0.456, 0.406])
imagenet_std = np.array([0.229, 0.224, 0.225])


def show_image(image, title=''):
    """
    显示图像，图像应该是一个 [H, W, 3] 的 numpy 数组，并且值已进行归一化。
    """
    # 确保图像是 RGB 格式
    assert image.shape[2] == 3
    # 将图像反归一化并显示
    plt.imshow(torch.clip((image * imagenet_std + imagenet_mean) * 255, 0, 255).int())  # 反归一化图像
    plt.title(title, fontsize=16)  # 设置标题
    plt.axis('off')  # 关闭坐标轴
    return


def prepare_model(chkpt_dir, arch='mae_vit_large_patch16'):
    """
    加载并准备 MAE 模型。

    参数：
    - chkpt_dir: 模型检查点路径
    - arch: 模型架构名称，默认为 'mae_vit_large_patch16'

    返回：
    - model: 加载了权重的 MAE 模型
    """
    # 根据架构名称构建模型
    model = getattr(models_mae, arch)()
    # 从检查点加载模型权重
    checkpoint = torch.load(chkpt_dir, map_location='cpu')
    # 将权重加载到模型中，并打印加载信息
    msg = model.load_state_dict(checkpoint['model'], strict=False)
    print(msg)  # 打印模型加载的状态
    return model


def run_one_image(img, model):
    """
    对单张图像进行 MAE 重建，展示原始图像、被遮挡的图像、重建图像以及重建与可见部分拼接后的图像。

    参数：
    - img: 输入的图像，已经是预处理后的 numpy 数组
    - model: 已加载的 MAE 模型
    """
    # 将输入图像转换为 tensor
    x = torch.tensor(img)

    # 扩展维度使其符合 batch 格式
    x = x.unsqueeze(dim=0)
    # 将图像的维度从 [batch, height, width, channels] 转换为 [batch, channels, height, width]
    x = torch.einsum('nhwc->nchw', x)

    # 使用 MAE 模型进行图像重建
    loss, y, mask = model(x.float(), mask_ratio=0.75)
    # 将模型输出的图像转换回原始图像形状
    y = model.unpatchify(y)
    # 将通道维度调整回 [batch, height, width, channels] 并转换为 CPU
    y = torch.einsum('nchw->nhwc', y).detach().cpu()

    # 可视化 mask
    mask = mask.detach()
    # 将 mask 转换为 [batch, H*W, p*p*3] 的形状
    mask = mask.unsqueeze(-1).repeat(1, 1, model.patch_embed.patch_size[0] ** 2 * 3)  # (N, H*W, p*p*3)
    mask = model.unpatchify(mask)  # 1 为移除，0 为保留
    mask = torch.einsum('nchw->nhwc', mask).detach().cpu()

    # 将输入图像的维度转换回 [batch, height, width, channels]
    x = torch.einsum('nchw->nhwc', x)

    # 被遮挡的图像：x * (1 - mask)
    im_masked = x * (1 - mask)

    # MAE 重建图像：将遮挡的部分使用重建图像填充
    im_paste = x * (1 - mask) + y * mask

    # 设置显示图像的大小
    plt.rcParams['figure.figsize'] = [24, 24]

    # 显示原始图像
    plt.subplot(1, 4, 1)
    show_image(x[0], "original")

    # 显示被遮挡的图像
    plt.subplot(1, 4, 2)
    show_image(im_masked[0], "masked")

    # 显示重建图像
    plt.subplot(1, 4, 3)
    show_image(y[0], "reconstruction")

    # 显示重建图像和可见部分拼接后的图像
    plt.subplot(1, 4, 4)
    show_image(im_paste[0], "reconstruction + visible")

    # 显示所有图像
    plt.show()


# 加载和预处理图像
img = Image.open(r'./0 五张代表性的/060-26P.jpg')  # 读取图像
img = img.resize((224, 224))  # 调整图像大小为 224x224
img = np.array(img) / 255.  # 将图像转为 numpy 数组并归一化到 [0, 1] 范围
assert img.shape == (224, 224, 3)  # 确保图像的形状为 (224, 224, 3)
# 使用 ImageNet 的均值和标准差对图像进行标准化
img = img - imagenet_mean
img = img / imagenet_std

# 设置显示图像的大小
plt.rcParams['figure.figsize'] = [5, 5]
# show_image(torch.tensor(img))  # 可选择显示图像

# 模型检查点路径
chkpt_dir = r'E:\CY\mae-main\mae-main\output_dir\checkpoint-399.pth'
# 加载 MAE 模型（mae_vit_base_patch16 架构）
model_mae = prepare_model(chkpt_dir, 'mae_vit_base_patch16')
print('Model loaded.')

# 设置随机种子以保证每次运行时的 mask 一致
torch.manual_seed(2)
print('MAE with pixel reconstruction:')
# 对图像进行处理并显示
run_one_image(img, model_mae)
